import os
import cv2 as cv
import numpy as np
import time


def auto_tune0(img_f, type_size, percent_f=0.001):
    row, col = img_f.shape
    # 将图片变为一维数组
    img_one = img_f.flatten()
    # 将一维数组排序
    img_one_sorted = np.sort(img_one)
    # 找出色调的上下限
    if percent_f == 0:
        low_cut = min(img_one_sorted)
        high_cut = max(img_one_sorted)
    else:
        low_cut = img_one_sorted[int(row * col * percent_f)]
        high_cut = img_one_sorted[int(row * col * (1 - percent_f))]
    # print('low,high:', low_cut, high_cut)
    # 对图片色调
    img_f[img_f < low_cut] = 0
    img_f[img_f > high_cut] = type_size
    img_f[(low_cut <= img_f) & (img_f <= high_cut)] = (type_size / (high_cut - low_cut)) * (img_f[(low_cut <= img_f) & (img_f <= high_cut)] - low_cut)
    return img_f


def auto_tone_one(img0):
    if img0.dtype == 'uint16':
        img_bit = 2 ** 16 - 1
    elif img0.dtype == 'uint8':
        img_bit = 2**8 - 1
    else:
        img_bit = None
        print('图片类型为：', img0.dtype, '，请手动添加img_bit')
    # 色调比例
    percent = 0.001
    # 自动色调
    img_out0 = auto_tune0(img0, img_bit, percent)
    # 合并为3通道
    # img_out1 = cv.merge((img_out0, img_out0, img_out0))
    return img_out0


def auto_tone_ones():
    # path[0]输入图片文件夹，path[1]为输出图片文件夹
    path = ['/diskb/ktb_set_new/Images/1/',
            '/diskb/ktb_set_new/1/']
    # path = ['/data/v5/test/1111/',
    #         '/data/v5/test/11100/']
    if not os.path.exists(path[1]):
        os.mkdir(path[1])
    names_list = os.listdir(path[0])
    for k in range(len(names_list)):
        print('正在处理处理%d/%d张图' % (k + 1, len(names_list)), names_list[k])
        start_time = time.time()
        # 读入图片
        img = cv.imread(os.path.join(path[0], names_list[k]), 0)
        print(img.shape)
        print('原图类型:', img.dtype)
        # 判断图片为多少位表示的颜色空间
        if img.dtype == 'uint16':
            img_bit = 2 ** 16 - 1
        else:
            img_bit = 255
        # 色调比例
        percent = 0.001
        # 自动色调
        img_out = auto_tune0(img, img_bit, percent)
        total_time = time.time() - start_time
        print('色调后图片类型：', img_out.dtype)
        print('用时%f s' % total_time)
        # 保存图片
        cv.imwrite(os.path.join(path[1], names_list[k]), img_out)


def auto_tone_color_ones():
    # path[0]输入图片文件夹，path[1]为输出图片文件夹
    path = ['/data1/TZB/ro_yolov5/convertor/org_img/',
            '/data1/TZB/ro_yolov5/convertor/a_img/']

    if not os.path.exists(path[1]):
        os.mkdir(path[1])
    names_list = os.listdir(path[0])
    for k in range(len(names_list)):
        name = names_list[k]
        print(name)
        print('正在处理处理%d/%d张图' % (k + 1, len(names_list)), names_list[k])
        start_time = time.time()
        # 读入图片
        img = cv.imread(os.path.join(path[0], names_list[k]))
        # print('原图类型:', img.dtype, img.shape)
        img_bit = 255
        # 色调比例
        percent = 0.001
        # 自动色调
        b_img, g_img, r_img = cv.split(img)
        b_out = auto_tune0(b_img, img_bit, percent)
        g_out = auto_tune0(g_img, img_bit, percent)
        r_out = auto_tune0(r_img, img_bit, percent)
        img_out = cv.merge((b_out, g_out, r_out))
        total_time = time.time() - start_time
        print('色调后图片类型：', img_out.dtype, img_out.shape)
        print('用时%f s' % total_time)
        # 保存图片
        cv.imwrite(os.path.join(path[1], names_list[k]), img_out)


def auto_tone_color_one(img, percent=0.1):
        img_bit = 255
        # 自动色调
        b_img, g_img, r_img = cv.split(img)
        b_out = auto_tune0(b_img, img_bit, percent)
        g_out = auto_tune0(g_img, img_bit, percent)
        r_out = auto_tune0(r_img, img_bit, percent)
        img_out = cv.merge((b_out, g_out, r_out))
        return img_out
